# 连接远程服务器并执行shell命令，将命令的执行结果存储。
import logging
import time
from datetime import datetime

import paramiko

from scan.models import Account, Machine
from task.models import CommandHistory


class SSH(object):
    def __init__(self, InstanceId):
        m = Machine.objects.filter(InstanceId=InstanceId).first()
        a = Account.objects.filter(User='root', Machine=m).first()
        if not m.PublicIpAddresses:
            logging.error(f'该实例{InstanceId}无公网IP，请人工排查')
            assert Exception(f'该实例{InstanceId}无公网IP，请人工排查')
        self.hostname = m.PublicIpAddresses[0]
        port = 22
        username = 'root'

        self.client = paramiko.SSHClient()
        # 自动添加策略，保存服务器的主机名和密钥信息
        self.client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
        if a.Password:
            # 连接SSH服务端，以用户名和密码进行认证
            self.client.connect(hostname=self.hostname, port=port, username=username, password=a.Password)
        elif a.Private_key:
            # 连接SSH服务端，以用户名和密码进行认证
            self.client.connect(hostname=self.hostname, port=port, username=username, password=a.Private_key)

    def do_cmd(self, cmd='hostname'):
        std2str = lambda x: x.read().decode('utf-8')
        start_time = time.time()
        stdin, stdout, stderr = self.client.exec_command(cmd)
        end_time = time.time()
        exit_code = stdout.channel.recv_exit_status()
        data_cmd = {
            'cmdId': cmd + str(time.time()),
            'status': 'success' if exit_code == 0 else 'failed',
            'type': 'shell',
            'cmd': cmd,
            'cmdResult': std2str(stdout),
            'ExitCode': exit_code,
            'StartTime': datetime.fromtimestamp(start_time),
            'RunTime': end_time - start_time,
        }
        CommandHistory.objects.create(**data_cmd)
        return data_cmd

if __name__ == '__main__':
    client = SSH(InstanceId="i-2zeiosvehx9da7w8ooqu")
    result = client.do_cmd('hostname')
    print(result)
